'''
对接百度的手写识别API
'''

import base64
import urllib
import requests
import json
import glob
import os
import math
import re
import cv2
from openai import OpenAI
from utils.config_utils import load_config
from utils.math_utils import group_numbers, find_largest_group
from utils.re_utils import remove_special_symbols, remove_letter_digit_combinations, remove_specific_combinations, remove_number_page_combinations
from utils.http_utils import post_request
from utils.cv2_utils import img2base64
import time
from dianxinApiSent import WritingOpenCorrectController
from utils.log_utils import SimpleLogger

simple_recog_logger = SimpleLogger('./log/recognition_pipeline.log')

API_KEY = "VAmp9Sw5YUNHLjDMF7Sb944i"
SECRET_KEY = "EPPIm8gbC4VAfcVNy35v1F1og23r3Jio"

config = load_config('config.json')
if config:
    API_KEY = config.get('hw_ak')
    SECRET_KEY = config.get('hw_sk')
    
def get_file_content_as_base64(path, urlencoded=False):
    """
    获取文件base64编码

    Parameters:
    - param path: 文件路径
    - param urlencoded: 是否对结果进行urlencoded 

    - return: base64编码信息
    """

    with open(path, "rb") as f:
        content = base64.b64encode(f.read()).decode("utf8")
        if urlencoded:
            content = urllib.parse.quote_plus(content)

    return content


def get_access_token():
    """
    使用 AK,SK 生成鉴权签名(Access Token)
    :return: access_token,或是None(如果错误)
    """
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    return str(requests.post(url, params=params).json().get("access_token"))


# 拼接两个文件路径，用来保存原始的识别结果和解析后的数据
def generate_save_path(file_path):

    original_data_path = None
    parsed_data_path = None
    file_extensions = ['.pdf', '.png', '.jpg']

    _, file_extension = os.path.splitext(file_path)
    if file_extension.lower() in file_extensions:
        original_data_path = file_path[0:-4] + '_original.txt'
        parsed_data_path = file_path[0:-4] + '_parsed.txt'

    return original_data_path, parsed_data_path

# 过滤一些非作文的文字，一般是答题卡上的文字或作答说明
def filter_irrelevant_words(words_result):

    try:
        filted_words = []
        for result in words_result:
            words = result['words']
            # 去除一些垂直方向的长文本获取高度比较小的文本（可能是行与行之间的字数提示）
            location = result['location']
            if location['height'] > location['width'] * 2 or location['height'] < 30:
                continue

            # TODO: 需要根据实际情况来修改过滤条件
            # 包含一些关键字则过滤掉，这个关键字可以持续增加
            irrelevant_words = ['答题', '作答', '答题卡', '姓名', '学号', "写作", ".作文", "座号", "0字", ")字", "提示线", "最低字"]
            contain_irrelevant_words = False
            for item in irrelevant_words:
                if words.find(item) >= 0:
                    contain_irrelevant_words = True
                    break
            
            # 不包含无关的词且宽度大于高度
            if not contain_irrelevant_words:
                remove_special_symbols_word = remove_special_symbols(words)
                remove_special_symbols_word = remove_number_page_combinations(remove_special_symbols_word)
                remove_special_symbols_word = remove_specific_combinations(remove_special_symbols_word)
                result['words'] = remove_special_symbols_word

                filted_words.append(result)

        return filted_words
    except Exception as e:
        print(e)
        return words_result
    

# 计算可能的分栏宽度
def calc_probable_column_width(words_result):

    witdhs_of_words = []
    for result in words_result:
        witdhs_of_words.append(result['location']['width'])

    # 对宽度进行分组
    grouped_widths = group_numbers(witdhs_of_words, threshold=200)

    # 找出数量最多的一组宽度即是最有可能的宽度
    max_group = find_largest_group(grouped_widths)
    #print('max_group: ', max_group)
    if len(max_group) <= 1:
        return None
    
    avg_width = sum(max_group) / len(max_group)
    #print('avg_width: ', avg_width)
    return avg_width

def group_words_by_left(words_result, threshold):
    """
    根据 left 属性对 locations 进行排序，并根据阈值进行分组。

    参数:
    - words_result (list of dict): 其中的location包含位置信息的字典列表，每个字典包含 top, left, width, height 属性。
    - threshold (int/float): 阈值，决定分组的距离。

    返回:
    - grouped_words_result (list of list of dict): 分组后的 locations 列表。
    """
    # 先根据 left 属性对 locations 进行排序
    sorted_words_result = sorted(words_result, key=lambda loc: loc['location']['left'])
    
    # 初始化分组结果
    grouped_words_result = []
    current_group = []

    # 遍历排序后的 locations
    for loc in sorted_words_result:
        if not current_group:
            # 如果当前分组为空，将第一个位置添加到当前分组
            current_group.append(loc)
        else:
            # 如果当前 location 与当前分组最后一个 location 的 left 差值小于等于阈值，将其添加到当前分组
            if math.fabs(loc['location']['left'] - current_group[-1]['location']['left']) <= threshold:
                current_group.append(loc)
            else:
                # 否则，当前分组结束，开始一个新的分组
                grouped_words_result.append(current_group)
                current_group = [loc]

    # 将最后一个分组添加到结果
    if current_group:
        grouped_words_result.append(current_group)

    return grouped_words_result

def sort_multiple_columns_words(words_result):

    thresh = 200
    probable_column_width = calc_probable_column_width(words_result)
    if probable_column_width is not None:
        thresh = round(probable_column_width / 2)

    # 根据每个location的left坐标进行排序，并根据某个阈值，对locations进行分组
    grouped_words_result = group_words_by_left(words_result, threshold=thresh)

    # 对每一组的location按top坐标排序，并合并起来
    sorted_words_result = []
    if len(grouped_words_result) > 1:
        for group in grouped_words_result:
            # TODO: 对每一组的location按top坐标排序
            sorted_group = sorted(group, key=lambda loc: loc['location']['top'])
            for item in sorted_group:
                sorted_words_result.append(item)
    else:
        sorted_words_result = words_result

    return sorted_words_result
    
def segment_paragraphs(ocr_data, indent_threshold=50):
    """
    根据OCR数据中的位置坐标分段落。

    参数:
    - ocr_data (list): 包含每行文本和位置的OCR数据列表。
    - indent_threshold (int): 识别段落的缩进阈值，默认值为50。

    返回:
    - list: 分段落的文本列表。
    """
    paragraphs = []
    current_paragraph = []

    for i, line in enumerate(ocr_data):
        if i == 0:
            paragraphs.append(line['words'])
            continue

        if i == 1:
            # 第一个段落的首行
            current_paragraph.append(line['words'])
        else:
            prev_left = ocr_data[i - 1]['location']['left']
            current_left = line['location']['left']

            # 如果当前行的left坐标比前一行的left坐标大，则认为是新段落
            if current_left > prev_left + indent_threshold:
                paragraphs.append("".join(current_paragraph))
                current_paragraph = [line['words']]
            else:
                current_paragraph.append(line['words'])

    # 最后一个段落
    if current_paragraph:
        paragraphs.append("".join(current_paragraph))

    return paragraphs

# 从识别结果中提取文字的内容，并尽量按照段落来拼接出完整的作文
def extract_words_from_recongnized_data(words_result):
    if words_result is None:
        return ''
    
    paragraphs = segment_paragraphs(words_result)
    parsed_text = ''
    for paragraph in paragraphs:
        parsed_text = parsed_text + paragraph + '\n\n'
    
    return parsed_text
  

# 请求识别接口并获取识别结果
def request_recognize_api(url, file_path, image_url=None):
    """
    请求识别接口并获取识别结果

    Parameters:
    - url (str): 识别接口的URL
    - file_path (str): 要识别的文件路径

    return:
    - dict or None: 识别接口的响应内容或在错误情况下返回None
    """
    begin_time = time.time()
    # 传图片的url
    if image_url is not None:
        payload = 'url=' + urllib.parse.quote_plus(image_url)
    else:
        # 传图片的路径
        if not os.path.exists(file_path):
            return None

        # 根据文件名称获取文件后缀名，并根据后缀名来生成请求的数据，目前仅支持image或pdf_file
        suffix = 'image'  # or pdf_file
        _, file_extension = os.path.splitext(file_path)
        if file_extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
            suffix = 'image'
        elif file_extension.lower() == '.pdf':
            suffix = 'pdf_file'
        else:
            # 不支持的文件格式
            return None
        
        file_as_base64 = get_file_content_as_base64(file_path,True)

        # 构造请求的payload
        payload = f'{suffix}={file_as_base64}'
    #print('payload: ', payload)
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': 'application/json'
    }    
    #print('payload : ', payload)
    response = requests.request("POST", url, headers=headers, data=payload)
    print('request_recognize_baidu_api time: ', time.time() - begin_time)
    # 根据response的status_code来决定返回response还是None
    if response.status_code == 200:
        return response.json()
    else:
        return None

# 识别和解析单份的pdf文件或图片文件，使用的是百度的识别接口
def proccesing_single_file(url, file_path, image_url=None, to_save=False):

    recognized_data = request_recognize_api(url, file_path, image_url=image_url)
    #print('proccesing_single_file recognized_data : ', recognized_data)

    parsed_text = None
    if recognized_data is not None and 'words_result' in recognized_data.keys():
        print('RECOGNITION SUCCESSFUL.')
        words_result = recognized_data['words_result']
        # 过滤掉一些非作文的文字
        words_result = filter_irrelevant_words(words_result)
        #print('words 1: ', words_result[0])
        # 对于有多栏切未分栏切割的情况，需要按照识别的坐标对识别结果进行分组排序
        sorted_words_result = sort_multiple_columns_words(words_result)
        #print('words 2: ', sorted_words_result[0])
        # 提取作文
        parsed_text = extract_words_from_recongnized_data(sorted_words_result)
        #print('words 3: ', parsed_text)
    
    if to_save:
        # 生成两个保存的路径，分别保存原始数据和提取的作文数据
        original_data_path, parsed_data_path = generate_save_path(file_path)
        # print(original_data_path)
        # print(parsed_data_path)
        if original_data_path is not None:
            # 保存原始识别数据
            with open(original_data_path, '+a', encoding='utf8') as fin:
                fin.write(json.dumps(recognized_data, ensure_ascii=False, indent=4))
                fin.close()
        else:
            print('ERROR SAVE PATH')

        if parsed_data_path is not None:
            # 保存提取的数据
            with open(parsed_data_path, '+w', encoding='utf8') as fin:
                fin.write(parsed_text)
                fin.close()
        else:
            print('ERROR SAVE PATH') 

    return parsed_text   

def recognize_handwriting_content(file, image_url=None, to_save=False):
    url = "https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token=" + get_access_token()
    return proccesing_single_file(url, file, image_url=image_url, to_save=to_save)


class HandWrittenRecognizer:
    def __init__(self, model_name:str) -> None:
        '''
        构造函数
        '''
        self.model_name = model_name
    
    def recognize(self, image_url, task):
        '''
        识别任务
        '''
        return ""
    
    def load_config(self, config_file:str):
        '''
        加载配置文件
        '''
        # self.config = load_config(config_file)
        pass

    def _load_config(self, config_file:str):
        '''
        加载配置文件
        '''
        try:
            with open(config_file, 'r', encoding='utf-8') as file:
                config = json.load(file)
                file.close()
            # return config.get('model', 'gpt-4')
            return config
        except Exception as e:
            print(e)
            return None
    

class BaiduHandWrittenRecognizer(HandWrittenRecognizer):
    '''
    使用百度的API进行手写识别
    '''
    def __init__(self, model_name:str, config_file:str) -> None:
        '''
        构造函数
        '''
        super().__init__(model_name)
        self.api_url = None
        self.api_key = None
        self.api_secretkey = None
        self.load_config(config_file)

    def recognize(self, image_url, task):
        '''
        识别任务
        '''
        print('config param : ', self.api_url, self.api_key, self.api_secretkey)
        return self._recognize(image_url, task)
    
    def _recognize(self, image_url, task):
        url = self.api_url + self._get_access_token()
        #print('merged url : ', url)
        return proccesing_single_file(url, None, image_url=image_url, to_save=False)

    def load_config(self, config_file:str):
        '''
        加载配置文件
        '''
        print('loading config file.')
        config = self._load_config(config_file)
        if config is not None:
            self.api_key = config.get('hw_ak')
            self.api_secretkey = config.get('hw_sk')
            self.api_url = config.get('baidu_hw_url')
        print('config param : ', self.api_url, self.api_key, self.api_secretkey)

    def _get_access_token(self):
        """
        使用 AK，SK 生成鉴权签名（Access Token）
        :return: access_token，或是None(如果错误)
        """
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {"grant_type": "client_credentials", "client_id": self.api_key, "client_secret": self.api_secretkey}
        return str(requests.post(url, params=params).json().get("access_token"))

class DianxinHandWrittenRecognizer(HandWrittenRecognizer):
    '''
    使用电信的API进行手写识别
    '''
    def __init__(self, model_name:str, config_file:str) -> None:
        '''
        构造函数
        '''
        super().__init__(model_name)
        self.recog_cap_no = "K1823621972597166080"
        self.load_config(config_file)
        self.controller = WritingOpenCorrectController(cap_no=self.recog_cap_no)

    def recognize(self, image_url, task):
        '''
        识别任务
        '''
        return self._recognize(image_url, task)

    def load_config(self, config_file:str):
        '''
        加载配置文件
        '''
        print('loading config file.')
        config = self._load_config(config_file)
        if config is not None:
            self.recog_cap_no = config.get('recog_cap_no')

    def _recognize(self, image_url, task):
        '''
        识别任务
        '''
        
        try:
            name = "test.png"
            original_result = self.controller.recognize(image_url, name, task)
            result = None
            if original_result is not None and original_result != "":
                result_json = json.loads(original_result)
                if "code" in result_json and "result" in result_json:
                    if result_json["code"] == 200:
                        result = result_json["result"]
            return result
        except Exception as e:
            print(e)
            return None


class DoubaoHandWrittenRecognizer(HandWrittenRecognizer):
    '''
    使用电信的API进行手写识别
    '''
    def __init__(self, config_file:str) -> None:
        '''
        构造函数
        '''
        super().__init__("doubao")
        self.api_url = "https://ark.cn-beijing.volces.com/api/v3"
        self.api_key = "bef29778-8a39-419a-8268-8dc874927cb8"
        self.prompt_template_for_doubao_formula_recog = None
        self.load_config(config_file)

    def recognize(self, image_url, task='formula', data_type='image_url'):
        '''
        识别任务

        task: formula or essay
        type: image_url or base64
        
        '''
        return self._recognize(image_url, task, data_type=data_type)

    def load_config(self, config_file:str):
        '''
        加载配置文件
        '''
        print('loading config file.')
        config = self._load_config(config_file)
        if config is not None:
            # self.api_url = config.get('baidu_predict_url')
            self.prompt_template_for_doubao_formula_recog = config.get('prompt_template_for_doubao_formula_recog', None)
        # print('config param : ', self.api_url)

    def _recognize(self, image_url, task='formula', data_type = 'image_url'):
        '''
        使用豆包进行手写识别
        image_url: 图片的url
        task: 识别的任务类型，目前支持formula和essay
        '''
        if image_url is None:
            return None
        simple_recog_logger.info('[START]recognizing image.')
        try:

            client = OpenAI(
                #api_key = os.environ.get("ARK_API_KEY"),
                api_key = self.api_key,
                base_url = self.api_url,
            )
    
            # 构造请求消息
            if data_type == 'base64':
                messages = [
                    {
                        # 指定消息的角色为用户
                        "role": "user",
                        "content": [
                            # 文本消息，希望模型根据图片信息回答的问题
                            {"type": "text", "text": self.prompt_template_for_doubao_formula_recog},
                            # 图片信息，希望模型理解的图片
                            {"type": "image_url", "image_url": {"url": "data:image/png;base64,"+image_url}},
                        ],
                    }
                ]
                # print("messages : ", messages)
            else:
                messages = [
                    {
                        # 指定消息的角色为用户
                        "role": "user",  
                        "content": [  
                            # 文本消息，希望模型根据图片信息回答的问题
                            {"type": "text", "text": self.prompt_template_for_doubao_formula_recog},  
                            # 图片信息，希望模型理解的图片
                            {"type": "image_url", "image_url": {"url":  image_url}},
                        ],
                    }
                ]

            # print("messages : ", messages)


            print("----- standard request -----")
            completion = client.chat.completions.create(
                # model = "ep-20250421145607-5zlkx",  # your model endpoint ID
                model = 'ep-20250908184955-qlck4',
                # model = 'ep-20250520192100-bv46c',
                messages = messages,
            )
            print('test.....')

            result = completion.choices[0].message.content
            simple_recog_logger.info('llm model result : {}'.format(repr(result)) )

            # print('doubao result : ', result, type(result))

            result_dict = None
            if isinstance(result, str):
                try:
                    # 尝试直接解析JSON
                    import demjson3 as demjson
                    result_dict = demjson.decode(result.strip())
                    # result_dict = json.loads(result.strip())
                # except json.JSONDecodeError as e:
                #     try:
                #         # 如果直接解析失败，尝试处理LaTeX中的反斜杠
                #         # 先将已经正确转义的序列临时替换
                #         temp_result = result.strip()
                #         # 替换已经存在的双反斜杠为临时标记
                #         temp_result = temp_result.replace('\\\\', '##ESCAPED_BACKSLASH##')
                #         # 替换单个反斜杠为双反斜杠
                #         temp_result = temp_result.replace('\\', '\\\\')
                #         # 恢复临时标记为双反斜杠
                #         temp_result = temp_result.replace('##ESCAPED_BACKSLASH##', '\\\\')
                        
                #         result_dict = json.loads(temp_result)
                #     except json.JSONDecodeError as e2:
                #         print(f"JSON解析错误: {e2}")
                #         simple_recog_logger.info(f'[END]recognizing image. with exception : {e2}')
                #         return None
                except Exception as e:
                    print(f"其他错误: {e}")
                    simple_recog_logger.info(f'[END]recognizing image. with exception : {e}')
                    return None
            else:
                result_dict = result
            

            simple_recog_logger.info('[END]recognizing image.')
            return result_dict
        except Exception as e:
            print(e)
            simple_recog_logger.info('[END]recognizing image. with exception : {}'.format(e) )
            return None
            


if __name__ == '__main__':

    # dir_path = "D:\\me\\zuowen\\2184609"
    # main(dir_path)

    # #file = "D:\\temp\\essay\\3.png"
    # file = "D:\\work\\code\\fx-business\\data\\8849324_0229201156_张旋_30.png"
    # url = "https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token=" + get_access_token()
    # print('url : ', url)
    # image_url = 'https://yj-oss.yunxiao.com/v1/baidu-raw/9017520/28297/4/000295/69475192-5BB5-11EF-8004-98EECBB5860E.png%40c_1%2Cx_1531%2Cy_333%2Cw_1349%2Ch_1057%7Cf_auto?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2024-08-22T03%3A03%3A29Z%2F3600%2Fhost%2Ffb22de2d871378686189aaf16b8249f84952acbe33cd271877929ac9d79b574c'
    # proccesing_single_file(url, file, image_url=None,  to_save=False)

    # exam_path = "D:\\me\\zuowen\\2184609"
    # pdf_filse = list_pdf_file(exam_path)
    # original_data_path, parsed_data_path = generate_save_path(pdf_filse[1])
    # print(original_data_path)
    # print(parsed_data_path)

    # image_url = "https://yj-oss.yunxiao.com/v1/baidu-raw/8950030/0/12/000090/2A772D77-3D9E-11EF-9DE0-FDA24EC8C8E1.png?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2024-08-14T07%3A21%3A19Z%2F3600%2Fhost%2Fb0047de96c2a91fcfd3e1dd2d0f94a27fe86bd2ac054502625fa826bd9a3e0fa"
    # url = "https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token=" + get_access_token()
    # text = proccesing_single_file(url, None, image_url=image_url, to_save=False)
    # print('text: ', text)

    # dianxin_recognizer = DianxinHandWrittenRecognizer(model_name='dianxin', config_file='config.json')
    # # dianxin_recognizer = BaiduHandWrittenRecognizer(model_name='ayx', config_file='config.json')
    # image_url = "https://yj-oss.yunxiao.com/v1/baidu-raw/8989575/12965/2/000214/0D0ACA55-3CE4-11EF-B316-6C4B9094B4C2.png?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2024-12-18T08%3A38%3A27Z%2F3600%2Fhost%2F1936916803088326b4cdba449c3d2a3d000101f2937f6869766a0ee969ffa8db"
    # task = 'cn'
    # recognized_text = dianxin_recognizer.recognize(image_url, task)
    # print('recognized_text : ', recognized_text)

    doubao_recognizer = DoubaoHandWrittenRecognizer(config_file='config.json')
    image_url = "https://yj-oss.yunxiao.com/v1/baidu-raw/9577873/107693/11/000012/AD6230B1-230B-11F0-A1D9-4C312DBB2C64.jpg%40c_1%2Cx_130%2Cy_239%2Cw_974%2Ch_255%7Cf_auto?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2025-04-27T06%3A57%3A21Z%2F3600%2Fhost%2F3e7a661983aac1c1f72bf1466ae72c231723e1f7c8c0e2ff4788059cb631fd8b"
    image_url = "https://yj-oss.yunxiao.com/v1/baidu-raw/9671959/35882/11/000023/C47D3670-40F9-11F0-A124-74DFBF4AE99E.jpg%40c_1%2Cx_768%2Cy_1356%2Cw_780%2Ch_99%7Cf_auto?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2025-06-04T14%3A21%3A44Z%2F3600%2Fhost%2Fc658b4f31ed685207da39a903879d922220faf3053ed0b1a542935ee5a02247c"
    task = 'formula'
    data_type = 'url' #'base64'
    # image_path = 'qwe.png'
    # # 读取图片并转换为base64编码
    # img = cv2.imread(image_path)
    # x, y, w, h = 679, 56, 250, 50
    # roi_img = img[y:y+h, x:x+w]
    # print(roi_img.shape)
    # base64_image = img2base64(roi_img)
    recognized_text = doubao_recognizer.recognize(image_url, task, data_type=data_type)
    # recognized_text = doubao_recognizer.recognize(image_url, task, data_type=data_type)
    print('recognized_text : ', recognized_text)
    

